package com.ask.forMe.langchain4j.rag.splitter;

import com.ask.forMe.langchain4j.model.entity.Section;
import com.ask.forMe.langchain4j.rag.loader.MyDocumentLoader;
import com.ask.forMe.langchain4j.rag.utils.StringSizeCalculator;
import com.ask.forMe.langchain4j.rag.word.WordProcessor;
import dev.langchain4j.data.segment.TextSegment;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 对markdown按标题进行分割
 */
@Slf4j
public class SplitMdByHeadingUtil {
    /**
     * 处理markdown文档
     *
     * @param filePath
     * @param maxSegmentSize
     * @param maxHeadingLevel
     * @return
     */
    public static List<Section> splitMDWithFullHeadingPath(String filePath, int maxSegmentSize, int maxHeadingLevel) {
        // 读取文档
        String text = MyDocumentLoader.loadAndParserByApacheTika(filePath).text();
        return basicAndSecondSplit(filePath, maxSegmentSize, maxHeadingLevel, text);
    }

    /**
     * 处理docx 文档
     *
     * @param filePath
     * @param maxSegmentSize
     * @param maxHeadingLevel
     * @return
     * @throws Exception
     */
    public static List<Section> splitDocxWithFullHeadingPath(String filePath, int maxSegmentSize, int maxHeadingLevel) throws Exception {
        // 读取文档1
        String text = WordProcessor.convertWordToMarkdownText(filePath);
        return basicAndSecondSplit(filePath, maxSegmentSize, maxHeadingLevel, text);
    }

    /**
     * 基础分割 + 2次分割
     *
     * @param filePath
     * @param maxSegmentSize
     * @param maxHeadingLevel
     * @return
     */

    private static List<Section> basicAndSecondSplit(String filePath, int maxSegmentSize, int maxHeadingLevel, String text) {
        log.info(
                "准备分割文档... 【文档路径】: {}, 【切分方式】: 标题, 【切分粒度】: {}, 【最大段长】: {}",
                filePath,
                maxHeadingLevel,
                maxSegmentSize
        );
        // 基础分割：基于标题，使用正则分割
        List<Section> basicSections = basicSplit(text, maxHeadingLevel);
        log.info("基础分割全部完成, 段数为: {}", basicSections.size());
        // 2次分割：对过长的分块2次分割，指定最大长度maxSegmentSize
        List<Section> secondSections = secondSplit(basicSections, maxSegmentSize);
        log.info("2次分割全部完成, 最终段数为: {}", secondSections.size());
        return secondSections;
    }


    /**
     * 对文档进行基础分割，并添加分块的标题路径
     *
     * @param text
     * @param maxHeadingLevel
     * @return
     */
    public static List<Section> basicSplit(String text, int maxHeadingLevel) {
        // 使用正则进行基于标题的分割
        String[] rawSections = RegexSplitUtils.splitByHeading(text, 1, maxHeadingLevel);
        List<Section> sections = new ArrayList<>();

        // 初始化每个分块的信息
        for (String content : rawSections) {
            if (!content.trim().isEmpty()) {
                int level = getHeadingLevel(content);
                if (level > 0) { // 有标题的部分
                    sections.add(new Section(content, level));
                } else { // 无标题部分（如开头内容）
                    sections.add(new Section(content, 0));
                }
            }
        }

        // 记录每个级别的最后一个分块
        Map<Integer, Section> lastSectionByLevel = new HashMap<>();

        // 从上到下处理标题层级
        for (Section section : sections) {
            if (section.getLevel() <= 1) { // 一级标题或无标题部分
                List<String> path = new ArrayList<>();
                if (section.getLevel() == 1) { // 当前分块为1级标题
                    path.add("H1-" + extractTitle(section.getContent())); // 只需添加1级标题路径，即自身标题
                }
                section.setTitlePath(path);
            } else {
                // 获取上一级标题信息
                Section parent = lastSectionByLevel.get(section.getLevel() - 1); // 查找最近的上级模块
                if (parent != null) {
                    // 复制上级标题路径
                    List<String> path = new ArrayList<>(parent.getTitlePath());
                    // 添加自己的标题
                    path.add("H" + section.getLevel() + "-" + extractTitle(section.getContent()));
                    section.setTitlePath(path);
                }
            }

            // 更新当前级别的最后一个分块
            lastSectionByLevel.put(section.getLevel(), section);

            // 去除内容中的标题信息
            section.setContent(removeFirstHeadingLine(section.getContent()));
        }
        return sections;
    }

    /**
     * 2次分割：对过长的分块进行2次分割
     *
     * @param maxSegmentSize
     * @param sections
     * @return
     */
    private static List<Section> secondSplit(List<Section> sections, int maxSegmentSize) {
        List<Section> result = new ArrayList<>();
        for (Section section : sections) {
            // 计算当前分块的长度
            StringSizeCalculator calculator = StringSizeCalculator.lengthCalculator();
            int l = calculator.sizeOf(section.getContent()); // 当前分块长度
            // 若某个分块的长度超过阈值，再按照段落分割
            if (l > maxSegmentSize) {
                List<TextSegment> segments = LangchainSplitter
                        .splitByParagraph(section.getContent(), maxSegmentSize, 0);
                List<Section> sectionList = segments
                        .stream()
                        .map(segment -> new Section(segment.text(), section.getLevel(), section.getTitlePath()))
                        .toList();
                result.addAll(sectionList);
                log.info(
                        "分块【{}】, 长度为: {}, 超过阈值: {}, 将其2次分割为: {}段...",
                        section.getTitlePath(),
                        l,
                        maxSegmentSize,
                        segments.size()
                );
            } else {
                // 不满足条件，保持原样
                result.add(section);
            }
        }
        return result;
    }

    /**
     * 从Markdown文本中提取标题级别
     *
     * @param text
     * @return
     */
    private static int getHeadingLevel(String text) {
        // 创建正则表达式模式，用于匹配行首的一个或多个#后跟至少一个空格
        Pattern pattern = Pattern.compile("^(#+)\\s+", Pattern.MULTILINE);
        // 在输入文本中查找匹配
        Matcher matcher = pattern.matcher(text);
        // 如果找到匹配项
        if (matcher.find()) {
            return matcher.group(1).length(); // 返回#符号的数量，即标题级别
        }
        return 0;
    }

    /**
     * 从Markdown文本中提取标题内容
     *
     * @param text
     * @return
     */
    private static String extractTitle(String text) {
        // 创建正则表达式模式，匹配行首的#号（一个或多个）、空格，然后捕获后面的所有内容
        Pattern pattern = Pattern.compile("^#+\\s+(.+)$", Pattern.MULTILINE);
        // 在输入文本中查找匹配
        Matcher matcher = pattern.matcher(text);
        if (matcher.find()) {
            return matcher.group(1).trim(); // 返回第一个捕获组（即标题内容部分），并去除前后空白
        }
        return "";
    }


    /**
     * 去除文本的第一行
     *
     * @param text
     * @return
     */
    public static String removeFirstHeadingLine(String text) {
        if (text == null || text.isEmpty()) return "";
        String[] lines = text.split("\n");
        String result = String.join("\n", Arrays.copyOfRange(lines, 1, lines.length));
        return result;
    }

}



